from langgraph.graph import StateGraph
from typing_extensions import TypedDict, Optional
from langgraph.graph import START, END
from langchain_openai import ChatOpenAI
from IPython.display import display,Image


class InputState(TypedDict):
  question: str
  llm_answer: Optional[str]

class OutputState(TypedDict):
  answer: str

class OverallState(InputState, OutputState):
  pass

complier = StateGraph(OverallState,input=InputState,output=OutputState)

def agent_node(state: InputState):
  # message = [{'role':'system', 'content':'你是一个智能体，你需要根据用户问题，生成一个回答'},
  #               {'role':'user', 'content':state['question']}]
  message=[('system', '你是一名知识渊博的专家，你需要根据用户问题，生成一个言简意赅的回答。'),
                ('human', state['question'])]
  llm = ChatOpenAI(
    base_url='http://localhost:11434/v1',
    model='gemma3:1b',
    api_key='zyy',
    
  )
  response = llm.invoke(message)
  print('agent: ', response.content)
  return {'llm_answer':response.content}


def action_node(state: InputState):
  message=[('system', '你是一个翻译助手， 你负责将用户输入翻译成英语'),
                ('human', state['llm_answer'])]
  llm = ChatOpenAI(
    base_url='http://localhost:11434/v1',
    model='gemma3:1b',
    api_key='zyy',
    
  )
  response = llm.invoke(message)
  return {'answer':response.content}

complier.add_node('Qwen3',agent_node)
complier.add_node('action',action_node)
complier.add_edge(START,'Qwen3')
complier.add_edge('Qwen3','action')
complier.add_edge('action',END)
graph = complier.compile()
output = graph.invoke({'question':'天上有多少颗星星'})
print('End: ', output['answer'])
# 保存图片到文件
# graph_image = graph.get_graph(xray=True)
# png_data = graph_image.draw_mermaid_png()

# # 将二进制数据写入文件
# with open("graph.png", "wb") as f:
#     f.write(png_data)

# print("图片已保存为 graph.png")


